﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Management;
using System.Diagnostics;
using System.Runtime.InteropServices;

namespace PowerTriggersWcf
{
    internal class PWTProcess : IDisposable
    {
        [DllImport("kernel32.dll", SetLastError = true, CallingConvention = CallingConvention.Winapi)]
        [return: MarshalAs(UnmanagedType.Bool)]
        private static extern bool IsWow64Process([In] IntPtr hProcess, [Out] out bool lpSystemInfo);

        private static volatile bool hasTestedForWow64;
        private static volatile bool isWow64;
        private static volatile string systemRoot;

        private ManagementObject _managementObject;
        private Dictionary<string, object> _cachedValues;

        private PWTProcess()
        { }

        internal PWTProcess(ManagementObject managementObject)
        {
            _cachedValues = new Dictionary<string, object>();
            _managementObject = managementObject;
        }

        internal bool HasPermission
        {
            get
            {
                bool hasPermission;
                var exeProp = _managementObject.Properties["ExecutablePath"];
                if ((exeProp != null && exeProp.Value != null))
                {
                    hasPermission = true;
                }
                else
                {
                    var commandLineProp = _managementObject.Properties["CommandLine"];
                    hasPermission = (commandLineProp != null && commandLineProp.Value != null);
                }
                if (hasPermission)
                {
                    if (!CurrentProcessIsWow64())
                    {
                        return true;
                    }
                    
                    string exePath = this.ExePath;
                    if (systemRoot == null)
                    {
                        systemRoot = Environment.GetEnvironmentVariable("SystemRoot").ToUpper();
                    }
                    if (!exePath.ToUpper().StartsWith(systemRoot))
                    {
                        return true;
                    }
                }

                return false;
            }
        }

        // This method returns True if the current process is running on Wow, False otherwise
        // it will cache the result after the first call, this is because the method is slow to execute and the result will always be the same
        internal static bool CurrentProcessIsWow64()
        {
            if (!hasTestedForWow64) 
            {
                using (var process = System.Diagnostics.Process.GetCurrentProcess())
                {
                    bool outValue;
                    if (!(IsWow64Process(process.Handle, out outValue)))
                    {
                        //throw new Exception("IsWow64Process() failed");
                        EventLogHelper.LogError("Error occured in CurrentProcessIsWow64");
                        // Need to guess it is true of false.
                        // It is better to assume that the current process is not Wow64 as this is the most likely scenario.
                        // If the process is indeed Wow64 then the service will fail in certain areas.
                        isWow64 = false; 
                    }
                    else
                    {
                        isWow64 = outValue;
                    }
                    hasTestedForWow64 = true;
                }
            }

            return isWow64;
        }

        internal int ProcessId
        {
            get
            {
                return (int)(uint)this.GetPropertyValue("ProcessId");
            }
        }

        internal int SessionId
        {
            get
            {
                return (int)(uint)this.GetPropertyValue("SessionId");
            }
        }

        internal string ProcessName
        {
            get
            {
                return (string)this.GetPropertyValue("Name");
            }
        }

        internal string ExePath
        {
            get
            {
                if (!_cachedValues.ContainsKey("ExePath"))
                {
                    _cachedValues.Add("ExePath", GetProcessExePath(_managementObject));
                }

                return (string)_cachedValues["ExePath"];
            }
        }

        internal string Description
        {
            get
            {
                if (!_cachedValues.ContainsKey("Description"))
                {
                    FileVersionInfo fvi = ProcessLauncher.GetCachedFileVersionInfo(this.ExePath.ToUpper());
                    if (fvi == null)
                    {
                        _cachedValues.Add("Description", System.IO.Path.GetFileName(this.ExePath));
                    }
                    else
                    {
                        string desc = fvi.FileDescription;
                        string shortName = System.IO.Path.GetFileName(this.ExePath);
                        string description;
                        if (!string.IsNullOrEmpty(desc) && System.IO.Path.GetFileName(desc) != shortName)
                        {
                            description = string.Format("{0} ({1})", desc, shortName);
                        }
                        else
                        {
                            description = shortName;
                        }
                        _cachedValues.Add("Description", description);
                    }
                }

                return (string)_cachedValues["Description"];
            }
        }

        private object GetPropertyValue(string propertyName)
        {
            if (!_cachedValues.ContainsKey(propertyName))
            {
                _cachedValues.Add(propertyName, _managementObject.GetPropertyValue(propertyName));
            }

            return _cachedValues[propertyName];
        }

        #region IDisposable Members

        public void Dispose()
        {
            if (_managementObject != null)
            {
                _managementObject.Dispose();
                _managementObject = null;
            }
        }

        #endregion

        #region static methods

        internal static IList<PWTProcess> GetProcesses()
        {
            List<PWTProcess> processes = new List<PWTProcess>();

            using (ManagementObjectSearcher searcher = new ManagementObjectSearcher("SELECT ProcessId, SessionId, ExecutablePath, CommandLine, Name, Description from Win32_Process"))
            {
                foreach (ManagementObject process in searcher.Get())
                {
                    if ((uint)process.GetPropertyValue("ProcessId") > 0) // processid zero is 'Idle'
                    {
                        PWTProcess proc = new PWTProcess(process);
                        if (proc.HasPermission)
                        {
                            processes.Add(proc);
                        }
                        else
                        {
                            process.Dispose();
                        }
                    }
                    else
                    {
                        process.Dispose();
                    }
                }
            }

            processes.TrimExcess();

            return processes;
        }

        internal static IList<PWTProcess> GetProcesses(string processName)
        {
            return GetProcesses(processName, null);
        }

        internal static IList<PWTProcess> GetProcesses(string processName, IList<int> sessionIds)
        {
            StringBuilder sessionCondition = new StringBuilder();

            if (sessionIds != null && sessionIds.Count > 0)
            {
                foreach (int sessionId in sessionIds)
                {
                    if (sessionCondition.Length == 0)
                    {
                        sessionCondition.Append(" AND (");
                    }
                    else
                    {
                        sessionCondition.Append(" OR ");
                    }
                    sessionCondition.AppendFormat("SessionId = {0}", sessionId);
                }
                sessionCondition.Append(")");
            }

            List<PWTProcess> processes = new List<PWTProcess>();

            using (ManagementObjectSearcher searcher = new ManagementObjectSearcher("SELECT ProcessId, SessionId, ExecutablePath, CommandLine, Name, Description from Win32_Process Where Name = '" + processName + "'" + sessionCondition.ToString()))
            {
                foreach (ManagementObject process in searcher.Get())
                {
                    if ((uint)process.GetPropertyValue("ProcessId") > 0) // processid zero is 'Idle'
                    {
                        PWTProcess proc = new PWTProcess(process);
                        if (proc.HasPermission)
                        {
                            processes.Add(proc);
                        }
                        else
                        {
                            process.Dispose();
                        }
                    }
                    else
                    {
                        process.Dispose();
                    }
                }
            }

            processes.TrimExcess();

            return processes;
        }

        //if (include)
        //{
        //    if (validUsers != null)
        //    {
        //        if (!validUsers.ContainsKey(process.SessionId))
        //        {
        //            include = false;
        //        }
        //    }
        //}


        internal static string GetProcessExePath(ManagementObject process)
        {
            string processFilePath;

            if (process.Properties["ExecutablePath"] != null && process.Properties["ExecutablePath"].Value != null)
            {
                processFilePath = (string)process.GetPropertyValue("ExecutablePath");
            }
            else if (process.Properties["CommandLine"] != null && process.Properties["CommandLine"].Value != null)
            {
                processFilePath = (string)process.GetPropertyValue("CommandLine");
            }
            else
            {
                throw new System.ComponentModel.Win32Exception("Access Denied");
            }

            if (processFilePath.StartsWith("%"))
            {
                string envVarName = processFilePath.Substring(0, processFilePath.IndexOf("%", 1) + 1);
                string envVarValue = Environment.GetEnvironmentVariable(envVarName.Replace("%", string.Empty), EnvironmentVariableTarget.Machine);
                if (envVarValue != null)
                {
                    processFilePath = System.IO.Path.Combine(envVarValue, processFilePath.Substring(envVarName.Length));
                }
            }
            else if (processFilePath.ToLower().StartsWith(@"\systemroot\"))
            {
                const string envVarName = "%SystemRoot%";
                string envVarValue = Environment.GetEnvironmentVariable(envVarName.Replace("%", string.Empty));
                if (envVarValue != null)
                {
                    processFilePath = System.IO.Path.Combine(envVarValue, processFilePath.Substring(envVarName.Length));
                }
            }

            return System.IO.Path.GetFullPath(processFilePath);
        }

        internal static string GetProcessDescription(ManagementObject process)
        {
            string name = (string)process.GetPropertyValue("Name");
            string description = name;
            if (process.Properties["Description"] != null)
            {
                string desc = (string)process.GetPropertyValue("Description");
                if (!string.IsNullOrEmpty(desc) && System.IO.Path.GetFileName(desc) != name)
                {
                    description = string.Format("{0} ({1})", desc, name);
                }
            }

            return description;
        }
        #endregion
    }
}
